#include <glpk.h>
#include "graph.h"
#include "graphreader.h"
#include "randomgraphreader.h"
#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <vector>
#include <set>
#include <string>
#include <time.h>
#include <sys/time.h>
#include <fstream>
#include <math.h>
using namespace std;

// Comment out if we want original order:
//#define RANDOM_ORDER

// Comment out if we don't know m beforehand
#define KNOWN_M

void graphTest();
void graphreaderTest(string filename);
void randomgraphreaderTest(string filename);
void densestSubgraphTest();

double charikar2approximation(Graph g);
double densestSubgraphLP(Graph g);

timespec startTimer();
double elapsedTime(timespec start);

void kbinsDensestSubgraph(string graphfile, string outputfile, int k, int seed);
void naiveDensestSubgraph(string graphfile, string outputfile, int seed);

bool withProb(int a, int b); // returns true w.p. a/b, false otherwise

int main(int argc, char** argv)
{
    string graphfile = "../data/as20000102.txt";  // relatively small graph
    char temp[256];
    char now[16];
    sprintf(now, "%ld", time(NULL)); // for getting unique files each time this is run for a given graph
    string kbinsoutputfile, naiveoutputfile;
    int i, k, r, REPEATS = 10;

    if (argc < 4)
    {
	cout << "Usage: ./main graph-file output-prefix maxk [repeats]" << endl;
	return 0;
    }

    if (argc > 4)
	REPEATS = atoi(argv[4]);

    graphfile = string(argv[1]);
    
    for(k = 2; k <= atoi(argv[3]); ++k)
    {
	for (r = 1; r <= REPEATS; ++r)
	{
	    sprintf(temp, "%s-%s-%dBINS-%d-", argv[2], now, k, r);
	    kbinsoutputfile = string(temp);
	    kbinsDensestSubgraph(graphfile, kbinsoutputfile, k, r);
	}
    }

#ifdef RANDOM_ORDER    
    for (r = 1; r <= REPEATS; ++r)
    {
	sprintf(temp, "%s-%s-NAIVE-%d-", argv[2], now, r);
	naiveoutputfile = string(temp);
	naiveDensestSubgraph(graphfile, naiveoutputfile, r);
    }
#else
    sprintf(temp, "%s-%s-NAIVE-1-", argv[2], now);
    naiveoutputfile = string(temp);
    naiveDensestSubgraph(graphfile, naiveoutputfile, 0);
#endif

    return 0;
}

void kbinsDensestSubgraph(string graphfile, string outputfile, int k, int seed)
{
    int node1, node2, m, krootm, i, j;
    timespec start;
    double rho_H[k], rho;
    Graph H[k];
    ofstream timefile, densityfile;
    double sumTime = 0.0, runTime;
    int edges = 0, totalEdges;
 
    timefile.open((outputfile + "TIME.txt").c_str());
    densityfile.open((outputfile + "DENSITY.txt").c_str());

#ifdef KNOWN_M
    GraphReader edgeCounter(graphfile);
    m = 0;
    while (edgeCounter.nextEdge(node1, node2)) { m++; }
#else
    m = 100;
#endif

    krootm = (int)pow(m, 1.0/k);

    for(i = 0; i < k; ++i)
	rho_H[i] = 0.0;

#ifdef RANDOM_ORDER
    RandomGraphReader gr(graphfile, seed);
#else
    GraphReader gr(graphfile);
#endif 

    // initialize seed after the random order graph code
    srand(seed);

    while (gr.nextEdge(node1, node2))
    {
	edges++;
	start = startTimer();

	H[0].addEdge(node1, node2);

	if (H[0].numEdges() < krootm)
	{
	    rho_H[0] = max(rho_H[0], charikar2approximation(H[0])); 
	}
	else
	{
	    // figure out what bin to move all the lower-edge bins to
	    i = 0;
	    totalEdges = H[0].numEdges();
	    while(i + 1 < k && totalEdges >= (int)pow(m, (i + 1.0)/k))
	    {
		i++;
		totalEdges += H[i].numEdges();
	    }

	    // move all the edges from bins 0, 1, ..., i-1 to bin i
	    for(j = 0; j < i; ++j)
	    {
		// copy all edges from H[j] to H[i] ...
		set<int> nodes_H(H[j].getNodes());
		for (set<int>::iterator n = nodes_H.begin(); n != nodes_H.end(); ++n)
		{
		    int node = *n;
		    set<int> nbrs = H[j].getNeighbors(node);
		    for (set<int>::iterator n2 = nbrs.begin(); n2 != nbrs.end(); ++n2)
		    {
			int nbr = *n2;
			if (node <= nbr)  // only add each edge once
			    H[i].addEdge(node, nbr);
		    }
		}
		
		// ... and reset H[j]
		H[j].reset();
		rho_H[j] = 0.0;
	    }

	    // ... and re-compute H[i]'s density
	    rho_H[i] = max(rho_H[i], charikar2approximation(H[i])); 	    
	}

	// update m and sqrtm if necessary
	if (edges > m)
	{
	    m = m * 2;   // double m each time number of edges goes over
	    krootm = (int)pow(m, 1.0/k);
	}
   
	// calculate max density as the densest of the rho_H's
	rho = rho_H[0];
	for(i = 1; i < k; ++i)
	    rho = max(rho, rho_H[i]);

	runTime = elapsedTime(start);
	sumTime += runTime;
	//cout << rho << " " << runTime << endl;
	timefile << runTime << endl;
        densityfile << rho << endl;
    }

    timefile.close();
    densityfile.close();
    cout << k << "binsDensestSubgraph total time " << sumTime << " seconds." << endl;
}

void naiveDensestSubgraph(string graphfile, string outputfile, int seed)
{
    int node1, node2;
    timespec start;
    double density;
    Graph G;
    ofstream timefile, densityfile;
    double sumTime = 0.0, runTime;
    int iteration = 0;

#ifdef RANDOM_ORDER
    RandomGraphReader gr(graphfile, seed);
#else
    GraphReader gr(graphfile);
#endif

    timefile.open((outputfile + "TIME.txt").c_str());
    densityfile.open((outputfile + "DENSITY.txt").c_str());

    while (gr.nextEdge(node1, node2))
    {
	G.addEdge(node1, node2);

	// only compute the value every 10th iteration for speed
	//if (iteration % 10 == 0)
	
	    start = startTimer();
	    density = charikar2approximation(G);
	    runTime = elapsedTime(start);
	

	iteration++;
	sumTime += runTime;
	//cout << density << " " << runTime << endl;	
	timefile << runTime << endl;  // use the value from the last iteration it was measured
	densityfile << density << endl;
    }

    timefile.close();
    densityfile.close();
    cout << "naiveDensestSubgraph total time " << sumTime << " seconds." << endl;
}


int old_main(int argc, char** argv)
{
  time_t begin, end;
  string filename = "../data/soc-Epinions1.txt";
  if (argc > 1)
    filename = string(argv[1]);
  
  //graphTest();
  //densestSubgraphTest();

  // NOTE: The non-default constructor calls the undirected graph reading 
  //       method that assumes that each edge appears twice (both ways).

  time(&begin);
  Graph graph(filename);
  cout << filename << " has " << graph.numNodes() << " nodes and " << graph.numEdges() << " edges.\n";
  time(&end);
  cout << "Loading the graph took " << difftime(end, begin) << " seconds" << endl << endl;

  time(&begin);
  cout << "Two-approximation: " << charikar2approximation(graph) << endl;
  time(&end);
  cout << "took " << difftime(end, begin) << " seconds" << endl << endl;

  // NOTE: After change to pass by reference (for memory reasons), the graph will 
  //       have been changed (actually, emptied) so we have to re-load it anew here.
  graph.reset();
  graph.readGraph(filename);

  time(&begin);
  cout << "Exact: " << densestSubgraphLP(graph) << endl;
  time(&end);
  cout << "took " << difftime(end, begin) << " seconds" << endl << endl;

  
  //cout << filename << " has " << graph.numNodes() << " nodes and " << graph.numEdges() << " edges.\n"; // graph unchanged 
   
  return 0;
}


/* Timer functions */

timespec startTimer()
{
    timespec start;
    clock_gettime(CLOCK_REALTIME, &start);
    return start;
}

double elapsedTime(timespec start)
{
    timespec end;
    clock_gettime(CLOCK_REALTIME, &end);
    return (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec)/1000000000.0;
}


/* Implementation of Charikar's linear-time 2-approximation algorithm for densest subgraph.
   this function only returns the approximated density, not the subgraph itself.
   note: This is currently not linear-time as there is an additional log(maxdegree) factor for moving neighbors each iteration. 
   note: Graph g is passed using (default) copy constructor rather than pointer/reference; this is important as local copy gets destroyed in this algorithm.
*/
double charikar2approximation(Graph g)
{
  double best = ((double)g.numEdges())/g.numNodes();
  
  // sets of vertices by degree
  vector<set<int> > degreeSet(g.numNodes() + 1, set<int>());
  set<int> nodes = g.getNodes();
  int minDegree = g.numNodes(); // keep track of smallest degree node
  for (set<int>::iterator n = nodes.begin(); n != nodes.end(); ++n)
  {
    int node = *n;
    int degree = g.getDegree(node);
    degreeSet[degree].insert(node);
    if (degree < minDegree) minDegree = degree;
  }

  while (g.numNodes() > 1)
  {
    if (minDegree > 0)
      minDegree--; // minDegree may have gone down 1 from the last iteration
    while(degreeSet[minDegree].empty())
    {
      minDegree++;
    }
    
    // just pick the first element of this set
    int lowDegreeNode = *degreeSet[minDegree].begin();
    
    // update degree set
    set<int> nbrs = g.getNeighbors(lowDegreeNode);
    for(set<int>::iterator n = nbrs.begin(); n != nbrs.end(); ++n)
    {
      int nbr = *n;
      int degree = g.getDegree(nbr);
      // move neighbor one degreeSet down
      if (nbr != lowDegreeNode) // don't update if this is a self-loop edge
      {
	  degreeSet[degree].erase(nbr);
	  degreeSet[degree - 1].insert(nbr);
      }
    }
    degreeSet[g.getDegree(lowDegreeNode)].erase(lowDegreeNode); // remove this node itself 
    
    // remove this node from the graph
    g.removeNode(lowDegreeNode);
    
    // update density (TODO?: change to rational representation for precision?)
    double density = ((double)g.numEdges())/g.numNodes();
    if (density > best)
    {
      best = density;
      //cout << best << " (" << g.numEdges() << "/" << g.numNodes() << ")" << endl;
    }
  }
  
  return best;
}

/* Implementation of Charikar's LP for computing exact densest subgraph
   this function only returns the approximated density, not the subgraph itself
   
   The LP is defined as follows:
   maximize \sum_{ij} x_{ij}
   e_{ij} = x_{ij} - y_i,
   f_{ij} = x_{ij} - y_j, for all ij \in E
   g = \sum_i y_i 
   e_{ij} <= 0, f_{ij} <= 0, g <= 1, 0 <= x_{ij}, 0 <= y_i

   Note: Assumes that nodes are called 1, 2, ..., g.numNodes(); CHECK!
*/ 
double densestSubgraphLP(Graph g)
{
  int m = g.numEdges(), n = g.numNodes();
  int ia[4 * m + n + 1], ja[4 * m + n + 1];
  int c, edgeindex;
  double ar[4 * m + n + 1], z;
  char buf[100];

  glp_prob *lp;
  lp = glp_create_prob();
  glp_set_prob_name(lp, "densest_subgraph");
  glp_set_obj_dir(lp, GLP_MAX);

  /* Create the rows */

  glp_add_rows(lp, m * 2 + 1);

  // create two rows for each edge
  for (int i = 1; i <= m; ++i)
  {
    sprintf(buf, "e%d", i);
    glp_set_row_name(lp, i, buf);
    glp_set_row_bnds(lp, i, GLP_UP, 0.0, 0.0);
  
    sprintf(buf, "f%d", i);
    glp_set_row_name(lp, i + m, buf);
    glp_set_row_bnds(lp, i + m, GLP_UP, 0.0, 0.0);
  }

  // create a row for g
  glp_set_row_name(lp, 2 * m + 1, "g");
  glp_set_row_bnds(lp, 2 * m + 1, GLP_UP, 0.0, 1.0);

  /* Create the columns */

  glp_add_cols(lp, m + n);

  // add the nodes first
  for(int i = 1; i <= n; ++i)
  {
    sprintf(buf, "y%d", i);
    glp_set_col_name(lp, i, buf);
    glp_set_col_bnds(lp, i, GLP_LO, 0.0, 0.0);
    glp_set_obj_coef(lp, 1, 0.0);  // doesn't show up in maximization target
  }

  // add the edges next
  edgeindex = 1;
  c = 1;
  set<int> nodes = g.getNodes();
  for (set<int>::iterator itr = nodes.begin(); itr != nodes.end(); ++itr)
  {
    int i = *itr; 
    set<int> nbrs = g.getNeighbors(i);
    for (set<int>::iterator itr2 = nbrs.begin(); itr2 != nbrs.end(); ++itr2)
    {
      int j = *itr2;
      if (i <= j) // only deal with each edge once; do allow self-loops
      {
	// add a column for this edge
	sprintf(buf, "x%d,%d", i, j);
	glp_set_col_name(lp, n + edgeindex, buf);
	glp_set_col_bnds(lp, n + edgeindex, GLP_LO, 0.0, 0.0);
	glp_set_obj_coef(lp, n + edgeindex, 1.0); // coeff. 1 in maximization target

	// update the matrix for this edge
	ia[c]   = edgeindex; ja[c]   = n+edgeindex; ar[c]   = 1.0;   // e_{ij} = x_{ij} 
	ia[c+1] = edgeindex; ja[c+1] = i;           ar[c+1] = -1.0;  //          - y_i

	ia[c+2] = edgeindex+m; ja[c+2] = n+edgeindex; ar[c+2] = 1.0;   // f_{ij} = x_{ij}
        ia[c+3] = edgeindex+m; ja[c+3] = j;           ar[c+3] = -1.0;  //          - y_j

	edgeindex++;
	c += 4;
      }
    }
  }
  
  // add n more entries for g = \sum_i y_i
  for(int i = 1; i <= n; ++i)
  {
    ia[c] = 2 * m + 1; ja[c] = i; ar[c] = 1.0; 
    c++;
  }

  /* Solve the LP */
  glp_load_matrix(lp, 4 * m + n, ia, ja, ar);
  glp_smcp parm;
  glp_init_smcp(&parm);
  parm.msg_lev=GLP_MSG_OFF; // turn off all message by glp_simplex
  glp_simplex(lp, &parm);
  z = glp_get_obj_val(lp);
  glp_delete_prob(lp);
  return z;
}

// returns true w.p. a/b, false otherwise
bool withProb(int a, int b)
{
    return rand() % b < a;
}


// lightly tests the graph class
void graphTest()
{
  Graph g;
  for (int i = 0; i < 5; ++i)
    g.addEdge(i, (i + 1) % 5);
  
  g.printGraph();
  cout << g.numEdges() << " edges, " << g.numNodes() << " nodes\n" << endl;
  
  g.removeEdge(0, 1);
  cout << "After removing edge (0, 1):\n";
  g.printGraph();
  cout << g.numEdges() << " edges, " << g.numNodes() << " nodes\n" << endl;
  
  g.removeNode(3);
  cout << "After removing node 3:\n";
  g.printGraph();
  cout << g.numEdges() << " edges, " << g.numNodes() << " nodes\n" << endl;
}

// lightly test the densest subgraph algorithms
void densestSubgraphTest()
{
  Graph g;

  // complete graph on 10 vertices; should be 4.5
  int n = 10;
  for(int i = 1; i <= n; ++i)
    for(int j = i + 1; j <= n; ++j)
      g.addEdge(i, j);
  cout << "Exact : " << densestSubgraphLP(g) << endl;
  cout << "2-appx: " << charikar2approximation(g) << endl;

  g.reset();

  // complete graph on even vertices of 200 nodes; should be 49.5
  n = 200;
  for (int i = 1; i <= n; ++i)
    g.addNode(i);
  for(int i = 2; i <= n; i += 2)
    for(int j = i + 2; j <= n; j += 2)
      g.addEdge(i, j);
  cout << "Exact : " << densestSubgraphLP(g) << endl;
  cout << "2-appx: " << charikar2approximation(g) << endl;

  g.reset();

  // K_{n,n,n} (complete tripartite) - density 2n/3, min deg n
  // K_{n+2} - density (2*n*n + (n+1)(n+2)/2)/(4n+2), min deg (n+1)

  n = 20;
  for (int i = 1; i <= 4 * n + 2; ++i)
    g.addNode(i);
  for(int i = 1; i <= n; ++i)
    for(int j = 1; j <= n; ++j)
    {
      g.addEdge(i, n + j);
      g.addEdge(n + i, 2 * n + j);
    }
  for(int i = 3 * n + 1; i <= 4 * n + 2; ++i)
    for(int j = i + 1; j <= 4 * n + 2; ++j)
      g.addEdge(i, j);
  cout << "Exact : " << densestSubgraphLP(g) << endl;
  cout << "2-appx: " << charikar2approximation(g) << endl;
}

void graphreaderTest(string filename)
{
    GraphReader gr(filename);
    int node1, node2;
    
    while(gr.nextEdge(node1, node2))
    {
	cout << node1 << "  " << node2 << endl;
    }
}

void randomgraphreaderTest(string filename)
{
    RandomGraphReader rgr(filename, 0);
    int node1, node2;

    while(rgr.nextEdge(node1, node2))
    {
        cout << node1 << "  " << node2 << endl;
    }
}
